#include <linux/arm-smccc.h>
#include <linux/delay.h>
#include <linux/errno.h>
#include <linux/fs.h>
#include <linux/init.h>
#include <linux/io.h>
#include <linux/kernel.h>
#include <linux/module.h>
#include <linux/of.h>
#include <linux/of_address.h>
#include <linux/of_device.h>
#include <linux/slab.h>
#include <linux/string.h>
#include <linux/sysfs.h>
#include <linux/uaccess.h>

#define CPU_ON_FUNCID      0xC4000003
#define RTOS_MEM_NODE_NAME "rtos@mem"  // 设备树节点名称

static ssize_t               sysfs_show(struct kobject* kobj, struct kobj_attribute* attr, char* buf);
static ssize_t               sysfs_store(struct kobject* kobj, struct kobj_attribute* attr, const char* buf, size_t count);
static struct kobj_attribute rtos_boot_sysfs_attribute = __ATTR(rtos_boot, 0660, sysfs_show, sysfs_store);

struct rtos_mem_info {
    phys_addr_t paddr;  // 物理地址
    size_t      size;   // 大小
};
struct rtos_mem_info rtos_mem = {0};

enum
{
    PSIC_METHOD_SMC = 0,
    PSIC_METHOD_HVC = 1,
};
static bool method_smcc = PSIC_METHOD_SMC;

static int off_core(int core_id)
{
    char         path[256];
    struct file* file;
    loff_t       pos = 0;
    ssize_t      len;
    const char*  offline_value = "0";

    if(core_id < 0 || core_id >= num_online_cpus())
    {
        pr_err("Num online CPUs: %d. CPU %d is invalid or already offline!\n", num_online_cpus(), core_id);
        return -EINVAL;
    }

    snprintf(path, sizeof(path), "/sys/devices/system/cpu/cpu%d/online", core_id);

    file = filp_open(path, O_WRONLY, 0);
    if(IS_ERR(file))
    {
        pr_err("Failed to open %s\n", path);
        return PTR_ERR(file);
    }

    len = kernel_write(file, offline_value, strlen(offline_value), &pos);
    if(len < 0)
    {
        pr_err("Failed to write to %s: %zd\n", path, len);
        filp_close(file, NULL);
        return len;
    }

    filp_close(file, NULL);
    pr_info("Successfully offline core %d!\n", core_id);

    return 0;
}

static int boot_rtos(int core_id)
{
    struct arm_smccc_res res = {0};

    off_core(core_id);

    if(method_smcc == PSIC_METHOD_SMC)
    {
        arm_smccc_smc(CPU_ON_FUNCID, core_id, rtos_mem.paddr, 0, 0, 0, 0, 0, &res);
    }
    else if(method_smcc == PSIC_METHOD_HVC)
    {
        arm_smccc_hvc(CPU_ON_FUNCID, core_id, rtos_mem.paddr, 0, 0, 0, 0, 0, &res);
    }
    else
    {
        pr_err("Psci method not support!\n");
        return -1;
    }

    pr_info("Psci method: %d, core_id: %d, res: 0x%lx\n", method_smcc, core_id, res.a0);

    return 0;
}

static int load_rtos_firmware(char* fw_path)
{
    struct file*  fw_file = NULL;
    mm_segment_t  old_fs;
    loff_t        pos = 0;
    size_t        fw_size;
    char*         fw_data = NULL;
    void __iomem* rtos_mem_vaddr;
    int           ret = 0;

    // 设置内核空间的文件操作权限
    old_fs = get_fs();
    set_fs(KERNEL_DS);

    // 打开固件文件
    fw_file = filp_open(fw_path, O_RDONLY, 0);
    if(IS_ERR(fw_file))
    {
        pr_err("Failed to open firmware file %s.\n", fw_path);
        ret = PTR_ERR(fw_file);
        goto out_unmap;
    }

    // 获取固件文件大小
    fw_size = i_size_read(fw_file->f_path.dentry->d_inode);
    if(fw_size > rtos_mem.size)
    {
        pr_err("Firmware size is larger than the reserved memory region!\n");
        ret = -EINVAL;
        goto out_close;
    }

    // 分配内存来存储固件数据
    fw_data = kmalloc(fw_size, GFP_KERNEL);
    if(!fw_data)
    {
        pr_err("Failed to allocate memory for firmware data!\n");
        ret = -ENOMEM;
        goto out_close;
    }

    // 读取固件文件数据
    ret = kernel_read(fw_file, fw_data, fw_size, &pos);
    if(ret < 0)
    {
        pr_err("Failed to read firmware file\n");
        goto out_free;
    }

    // 映射物理内存到虚拟地址空间
    rtos_mem_vaddr = ioremap(rtos_mem.paddr, rtos_mem.size);
    if(!rtos_mem_vaddr)
    {
        pr_err("Failed to ioremap memory region.\n");
        return -ENOMEM;
    }
    // 将固件数据拷贝到预留的物理内存区域
    memcpy_toio(rtos_mem_vaddr, fw_data, fw_size);
    iounmap(rtos_mem_vaddr);

    pr_info("Firmware %s loaded into memory at %pa\n", fw_path, &rtos_mem.paddr);

out_free:
    kfree(fw_data);
out_close:
    filp_close(fw_file, NULL);
out_unmap:
    set_fs(old_fs);
    return ret;
}

static ssize_t sysfs_show(struct kobject* kobj, struct kobj_attribute* attr, char* buf)
{
    return sprintf(buf, "Module function triggered\n");
}

static ssize_t sysfs_store(struct kobject* kobj, struct kobj_attribute* attr, const char* buf, size_t size)
{
    int         icnt     = 0;
    int         ret      = 0;
    const char* load_cmd = "load";
    const char* boot_cmd = "boot";
    const char* mem_cmd  = "mem";

    if(strstr(buf, load_cmd) && strstr(buf, "-f"))
    {

        char  firmware_path[256] = {0};
        char* file_path_start    = strstr(buf, "-f") + 3;  // 跳过"-f "部分，并且必须跟一个空格。

        // TO DO: 检查file_path_start是否在buf+size范围内

        // 跳过空白字符
        while(*file_path_start == ' ' && file_path_start < (buf + size))
        {
            file_path_start++;
        }

        // 检查是否到buf末尾
        if(file_path_start >= (buf + size))
        {
            goto _invalid;
        }

        // 读取字符串，直到遇到空格或换行符，或者到达buf末尾
        while(*file_path_start != ' ' && *file_path_start != '\n' && file_path_start < (buf + size))
        {
            firmware_path[icnt++] = *file_path_start++;
        }

        load_rtos_firmware(firmware_path);
    }
    else if(strstr(buf, boot_cmd) && strstr(buf, "-c"))
    {
        int   core_id   = 0;
        char* arg_start = strstr(buf, "-c") + 3;

        while(*arg_start == ' ' && arg_start < (buf + size))
        {
            arg_start++;
        }

        if(arg_start >= (buf + size))
        {
            goto _invalid;
        }

        core_id = *arg_start - '0';

        boot_rtos(core_id);
    }
    else if(strstr(buf, mem_cmd))
    {
        if(strstr(buf, "-p") || strstr(buf, "-s"))
        {
            char*       arg_start;
            phys_addr_t paddr      = 0;
            size_t      paddr_size = 0;

            // TO DO: 检查参数格式

            // 配置运行地址
            if(strstr(buf, "-p"))
            {
                char parm_buf[32] = {0};
                long parm         = 0;

                arg_start = strstr(buf, "-p") + 3;

                arg_start = arg_start != NULL ? strstr(arg_start, "0x") : NULL;
                if(arg_start == NULL)
                {
                    arg_start = strstr(arg_start, "0X");
                    if(arg_start == NULL)
                    {
                        goto _invalid;
                    }
                }

                arg_start += 2;

                if(arg_start >= (buf + size))
                {
                    goto _invalid;
                }

                while(*arg_start != ' ' && *arg_start != '\n' && arg_start < (buf + size))
                {
                    parm_buf[icnt++] = *arg_start++;
                }
                parm_buf[icnt] = '\0';

                ret = kstrtol(parm_buf, 16, &parm);
                if(ret < 0)
                {
                    goto _invalid;
                }

                if(parm != 0)
                {
                    // rtos_mem.paddr = parm;
                    paddr = parm;
                }
                else
                {
                    goto _invalid;
                }
            }

            if(strstr(buf, "-s"))
            {
                char parm_buf[32] = {0};
                long parm         = 0;

                arg_start = strstr(buf, "-s") + 3;

                arg_start = arg_start != NULL ? strstr(arg_start, "0x") : NULL;
                if(arg_start == NULL)
                {
                    arg_start = strstr(arg_start, "0X");
                    if(arg_start == NULL)
                    {
                        goto _invalid;
                    }
                }

                arg_start += 2;

                if(arg_start >= (buf + size))
                {
                    goto _invalid;
                }

                while(*arg_start != ' ' && *arg_start != '\n' && arg_start < (buf + size))
                {
                    parm_buf[icnt++] = *arg_start++;
                }
                parm_buf[icnt] = '\0';

                ret = kstrtol(parm_buf, 16, &parm);
                if(ret < 0)
                {
                    goto _invalid;
                }

                if(parm != 0)
                {
                    // rtos_mem.size = parm;
                    paddr_size = parm;
                }
                else
                {
                    goto _invalid;
                }
            }

            rtos_mem.paddr = paddr != 0 ? paddr : rtos_mem.paddr;
            rtos_mem.size  = paddr_size != 0 ? paddr_size : rtos_mem.size;

            pr_info("Set rtos use mem addr: %pa size: 0x%lx\n", &rtos_mem.paddr, rtos_mem.size);
        }
        else
        {
            goto _invalid;
        }
    }
    else
    {
    _invalid:
        pr_info("Invalid operation of [%s]\n", buf);
        return -EINVAL;
    }

    return size;
}

static int rtos_mem_init(char* mem_node_name)
{
    struct device_node* np;
    u32                 addr, size;
    int                 ret;
    char                node_name[128] = "/reserved-memory/";

    if(mem_node_name != NULL)
    {
        sprintf(node_name, "/reserved-memory/%s", mem_node_name);
    }
    else
    {
        sprintf(node_name, "/reserved-memory/%s", RTOS_MEM_NODE_NAME);
    }

    pr_info("RTOS reserved memory node name: %s.\n", node_name);

    // 查找设备树中预留的内存节点
    np = of_find_node_by_path(node_name);
    if(!np)
    {
        pr_err("Failed to find node %s in device tree!\n", node_name);
        return -ENODEV;
    }

    // 获取 'reg' 属性
    ret = of_property_read_u32_index(np, "reg", 1, &addr);  // 获取地址
    if(ret)
    {
        pr_err("Failed to read 'reg' property for address!\n");
        return -EINVAL;
    }

    ret = of_property_read_u32_index(np, "reg", 3, &size);  // 获取大小
    if(ret)
    {
        pr_err("Failed to read 'reg' property for size!\n");
        return -EINVAL;
    }

    pr_info("Reserved memory region found at 0x%08x, size 0x%08x.\n", addr, size);

    rtos_mem.paddr = addr;
    rtos_mem.size  = size;

    return 0;
}

static void get_psci_method(void)
{
    const char*         method;
    struct device_node* np;
    struct of_device_id psci_of_match[] = {
        {.compatible = "arm,psci"},
        {.compatible = "arm,psci-0.2"},
        {.compatible = "arm,psci-1.0"},
        {},
    };

    np = of_find_matching_node(NULL, psci_of_match);
    if(!np)
    {
        pr_err("No matching PSCI node found in device tree!\n");
        return;
    }

    if(!of_device_is_available(np))
    {
        pr_err("PSCI node found, but it is not available!\n");
        of_node_put(np);
        return;
    }

    if(of_property_read_string(np, "method", &method))
    {
        pr_err("Failed to read 'method' property from PSCI node!\n");
        of_node_put(np);
        return;
    }

    of_node_put(np);

    if(strcmp(method, "hvc") == 0)
    {
        method_smcc = PSIC_METHOD_HVC;
        pr_info("PSCI method: hvc!\n");
    }
    else if(strcmp(method, "smc") == 0)
    {
        method_smcc = PSIC_METHOD_SMC;
        pr_info("PSCI method: smc!\n");
    }
    else
    {
        pr_err("Invalid PSCI method: %s!\n", method);
        return;
    }
}

static int __init rtos_boot_init(void)
{
    int error;

    // 查找设备树中预留的内存节点
    rtos_mem_init(NULL);

    // 查找smcc驱动
    get_psci_method();

    error = sysfs_create_file(kernel_kobj, &rtos_boot_sysfs_attribute.attr);
    if(error)
    {
        pr_err("Failed to create sysfs file!\n");
    }

    return 0;
}

static void __exit rtos_boot_exit(void)
{
    sysfs_remove_file(kernel_kobj, &rtos_boot_sysfs_attribute.attr);
    pr_info("RTOS boot module remove!\n");
}

module_init(rtos_boot_init);
module_exit(rtos_boot_exit);
MODULE_LICENSE("GPL");
